import gymnasium as gym
import math
import random
from typing import Dict, List, Tuple

import matplotlib.pyplot as plt
from collections import namedtuple, deque
from itertools import count
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

from segment_tree import MinSegmentTree, SumSegmentTree
import pickle
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# if GPU is to be used
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Transition = namedtuple('Transition',
                        ('state', 'action', 'next_state', 'reward'))
import os

class ReplayMemory(object):

    def __init__(self, capacity):
        self.memory = deque([], maxlen=capacity)
        self.max_size = capacity

    def push(self, *args):
        """Save a transition"""
        self.memory.append(Transition(*args))

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)


class PrioritizedReplayBuffer(ReplayMemory):
    """Prioritized Replay buffer.

    Attributes:
        max_priority (float): max priority
        tree_ptr (int): next index of tree
        alpha (float): alpha parameter for prioritized replay buffer
        sum_tree (SumSegmentTree): sum tree for prior
        min_tree (MinSegmentTree): min tree for min prior to get max weight

    """

    def __init__(
            self,
            capacity,
            alpha: float = 0.6,
            beta: float = 0.4
    ):
        """Initialization."""
        assert alpha >= 0
        assert beta >= 0

        super(PrioritizedReplayBuffer, self).__init__(capacity)
        self.max_priority, self.tree_ptr = 1.0, 0
        self.alpha = alpha
        self.beta = beta

        # capacity must be positive and a power of 2.
        tree_capacity = 1
        while tree_capacity < self.max_size:
            tree_capacity *= 2

        self.sum_tree = SumSegmentTree(tree_capacity)
        self.min_tree = MinSegmentTree(tree_capacity)

    def push(
            self,
            *args
    ):
        """Store experience and priority."""
        super().push(*args)

        self.sum_tree[self.tree_ptr] = self.max_priority ** self.alpha
        self.min_tree[self.tree_ptr] = self.max_priority ** self.alpha
        self.tree_ptr = (self.tree_ptr + 1) % self.max_size

    def sample(self, batch_size, beta:float):
        """Sample a batch of experiences."""
        assert len(self) >= batch_size
        assert beta > 0

        indices = self._sample_proportional(batch_size)

        weights = np.array([self._calculate_weight(i, beta) for i in indices])

        return [self.memory[index] for index in indices], weights, indices

    def update_priorities(self, indices: List[int], priorities: np.ndarray):
        """Update priorities of sampled transitions."""
        assert len(indices) == len(priorities)

        for idx, priority in zip(indices, priorities):
            assert priority > 0
            assert 0 <= idx < len(self)

            self.sum_tree[idx] = priority ** self.alpha
            self.min_tree[idx] = priority ** self.alpha

            self.max_priority = max(self.max_priority, priority)

    def _sample_proportional(self, batch_size) -> List[int]:
        """Sample indices based on proportions."""
        indices = []
        p_total = self.sum_tree.sum(0, len(self) - 1)
        segment = p_total / batch_size

        for i in range(batch_size):
            a = segment * i
            b = segment * (i + 1)
            upperbound = random.uniform(a, b)
            idx = self.sum_tree.retrieve(upperbound)
            indices.append(idx)

        return indices

    def _calculate_weight(self, idx: int, beta: float):
        """Calculate the weight of the experience at idx."""
        # get max weight
        p_min = self.min_tree.min() / self.sum_tree.sum()
        max_weight = (p_min * len(self)) ** (-beta)

        # calculate weights
        p_sample = self.sum_tree[idx] / self.sum_tree.sum()
        weight = (p_sample * len(self)) ** (-beta)
        weight = weight / max_weight

        return weight


# Deuling Network
class DQN(nn.Module):

    def __init__(self, n_observations: int, n_actions: int, num_cells: int = 128):
        super(DQN, self).__init__()

        self.n_observations = n_observations
        self.n_actions = n_actions
        # 继续增加层数效果会变差
        self.features_layer = nn.Linear(n_observations, num_cells) # feature layer
        
        self.advantage_hidden_layer = nn.Linear(num_cells, num_cells) # advantages layer 1
        self.advantage_layer = nn.Linear(num_cells, n_actions) #  advantages layer 2
        
        self.value_hidden_layer = nn.Linear(num_cells, num_cells)
        self.value_layer = nn.Linear(num_cells, 1) # value layer 2

    # Called with either one element to determine next action, or a batch
    # during optimization. Returns tensor([[left0exp,right0exp]...]).
    def forward(self, x):
        # Dueling network relu和tanh对比
        x = F.relu(self.features_layer(x))
        advantage_x = F.relu(self.advantage_hidden_layer(x))
        advantage = self.advantage_layer(advantage_x)

        value_x = F.relu(self.value_hidden_layer(x))
        value = self.value_layer(value_x)

        q = value + advantage - advantage.mean(dim=-1, keepdim=True)

        return q
    
def select_action(state):
    # epsilon greedy action selection
    global steps_done
    sample = random.random()
    eps_threshold = EPS_END + (EPS_START - EPS_END) * \
        math.exp(-1. * steps_done / EPS_DECAY)
    steps_done += 1
    if sample > eps_threshold:
        with torch.no_grad():
            # t.max(1) will return the largest column value of each row.
            # second column on max result is index of where max element was
            # found, so we pick action with the larger expected reward.
            return policy_net(state).max(1).indices.view(1, 1), eps_threshold
    else:
        return torch.tensor([[env.action_space.sample()]], device=device, dtype=torch.long), eps_threshold


def plot_durations(show_result=False):
    plt.figure(1)
    durations_t = torch.tensor(episode_durations, dtype=torch.float)
    if show_result:
        plt.title('Result')
    else:
        plt.clf()
        plt.title('Training...')
    plt.xlabel('Episode')
    plt.ylabel('Score')
    plt.plot(durations_t.numpy())
    # Take 100 episode averages and plot them too
    if len(durations_t) >= 100:
        means = durations_t.unfold(0, 100, 1).mean(1).view(-1)
        means = torch.cat((torch.zeros(99), means))
        plt.plot(means.numpy())


def optimize_model():
    if len(memory) < BATCH_SIZE:
        return
    transitions, weights, indices = memory.sample(BATCH_SIZE, BETA)
    weights = torch.FloatTensor(
        weights.reshape(-1, 1)
    ).to(device)
    # Transpose the batch (see https://stackoverflow.com/a/19343/3343043 for
    # detailed explanation). This converts batch-array of Transitions
    # to Transition of batch-arrays.
    batch = Transition(*zip(*transitions))

    # Compute a mask of non-final states and concatenate the batch elements
    # (a final state would've been the one after which simulation ended)
    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
                                            batch.next_state)), device=device, dtype=torch.bool)
    non_final_next_states = torch.cat([s for s in batch.next_state
                                                if s is not None])
    state_batch = torch.cat(batch.state)
    action_batch = torch.cat(batch.action)
    reward_batch = torch.cat(batch.reward)

    # Compute Q(s_t, a) - the models computes Q(s_t), then we select the
    # columns of actions taken. These are the actions which would've been taken
    # for each batch state according to policy_net
    state_action_values = policy_net(state_batch).gather(1, action_batch)

    # Compute V(s_{t+1}) for all next states.
    # Expected values of actions for non_final_next_states are computed based
    # on the "older" target_net; selecting their best reward with max(1).values
    # This is merged based on the mask, such that we'll have either the expected
    # state value or 0 in case the state was final.
    next_state_values = torch.zeros(BATCH_SIZE, device=device)
    with torch.no_grad():
        # origin
        # next_state_values[non_final_mask] = target_net(non_final_next_states).max(1).values

        # double DQN
        next_state_values[non_final_mask] = target_net(non_final_next_states).gather(1, policy_net(non_final_next_states).argmax(dim=1, keepdim=True)).squeeze(1) # 降维度为n
    # Compute the expected Q values
    expected_state_action_values = (next_state_values * GAMMA) + reward_batch
    # Compute Huber loss
    elementwise_loss = F.smooth_l1_loss(input=state_action_values, target=expected_state_action_values.unsqueeze(1),
                                        reduction='none', beta=1.0)

    # PER: importance sampling before average
    loss = torch.mean(elementwise_loss * weights)

    # Optimize the models
    optimizer.zero_grad()
    loss.backward()
    # In-place gradient clipping
    torch.nn.utils.clip_grad_value_(policy_net.parameters(), 100)
    optimizer.step()

    # PER: update priorities
    loss_for_prior = elementwise_loss.detach().cpu().numpy()
    new_priorities = loss_for_prior + PER_EPS
    memory.update_priorities(indices, new_priorities)
    
    return loss.item()


def save_best_model(path):
    try:
        #gpu tensor format save
        torch.save(obj={
            'policy_net': policy_net.state_dict(),
            'target_net': target_net.state_dict(),
            'BATCH_SIZE':BATCH_SIZE,
            'GAMMA' : GAMMA,
            'EPS_START' : EPS_START,
            'EPS_END' : EPS_END,
            'EPS_DECAY' : EPS_DECAY, # total frames * 0.1
            'TAU' : TAU,
            'LR' : LR,
            # PER parameters
            'ALPHA' : ALPHA,
            'BETA' : BETA,
            'PER_EPS' : PER_EPS,
            'MEMORY_SIZE' : MEMORY_SIZE
        }, f=path)
        logger.info("save models success")
    except:
        logger.warning("save models failed.")

HOME = os.path.dirname(os.path.realpath(__file__))
print(HOME)
os.makedirs(f'{HOME}/models', exist_ok=True)
os.makedirs(f'{HOME}/record', exist_ok=True)
max_episode_steps = 500
env_name = "CartPole-v1"
env = gym.make(f"{env_name}")
save_model_path = os.path.join(HOME, f'models/best_{env_name}_model.pt')
recorder_path = os.path.join(HOME, f'record/{env_name}_score.pkl')
num_episodes = 500

# BATCH_SIZE is the number of transitions sampled from the replay buffer
# GAMMA is the discount factor as mentioned in the previous section
# EPS_START is the starting value of epsilon
# EPS_END is the final value of epsilon
# EPS_DECAY controls the rate of exponential decay of epsilon, higher means a slower decay
# TAU is the update rate of the target network
# LR is the learning rate of the ``AdamW`` optimizer
BATCH_SIZE = 128
GAMMA = 0.99
EPS_START = 0.9
EPS_END = 0.05
EPS_DECAY = 5_000 # total frames * 0.1
TAU = 0.005
LR = 1e-4
# PER parameters
ALPHA = 0.2
BETA = 0.6
PER_EPS = 1e-5
MEMORY_SIZE = 50_000

episode_durations = []
# Get number of actions from gym action space
n_actions = env.action_space.n
# Get the number of state observations
state, info = env.reset()
n_observations = len(state)
policy_net = DQN(n_observations, n_actions, num_cells=128).to(device)
target_net = DQN(n_observations, n_actions, num_cells=128).to(device)
target_net.load_state_dict(policy_net.state_dict())
optimizer = optim.Adam(policy_net.parameters(), lr=LR)
memory = PrioritizedReplayBuffer(MEMORY_SIZE, alpha=ALPHA, beta=BETA) # as big as possible

if __name__ == '__main__':

    steps_done = 0
    num_frames = num_episodes* max_episode_steps
    frame_idx = 0
    
    # 训练效果记录
    recorder = dict({
        'frame': [],
        'epoch': [],
        'score': [],
        'loss': [],
        'epsilons': [],
    })
    
    for i_episode in range(num_episodes):
        # Initialize the environment and get its state
        state, info = env.reset()
        state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)

        score = 0
        for t in count():
            # PER: increase beta
            frame_idx += 1
            fraction = min(frame_idx / num_frames, 1.0)
            BETA = BETA + fraction * (1.0 - BETA)
            
            action, eps_threshold = select_action(state)
            observation, reward, terminated, truncated, _ = env.step(action.item())
            reward = torch.tensor([reward], device=device)
            done = terminated or truncated
            score += reward
            if terminated:
                next_state = None
            else:
                next_state = torch.tensor(observation, dtype=torch.float32, device=device).unsqueeze(0)

            # Store the transition in memory
            memory.push(state, action, next_state, reward)

            # Move to the next state
            state = next_state

            # Perform one step of the optimization (on the policy network)
            loss = optimize_model()

            # Soft update of the target network's weights
            # θ′ ← τ θ + (1 −τ )θ′
            target_net_state_dict = target_net.state_dict()
            policy_net_state_dict = policy_net.state_dict()
            for key in policy_net_state_dict:
                target_net_state_dict[key] = policy_net_state_dict[key]*TAU + target_net_state_dict[key]*(1-TAU)
            target_net.load_state_dict(target_net_state_dict)

            if done:
                
                # save to recorder
                recorder['epoch'].append(i_episode)
                recorder['score'].append(score.squeeze(0).cpu().numpy())
                recorder['loss'].append(loss)
                recorder['epsilons'].append(eps_threshold)
                
                print(
                    f'episode is {i_episode}, score is {score.squeeze(0).cpu().numpy()}, eps is {eps_threshold}')
                episode_durations.append(score)
                # plot_durations()
                break
            
    # save to file for ploting
    with open(recorder_path, 'wb') as f:
        pickle.dump(recorder, f)
    save_best_model(save_model_path)
    print('Complete')
